package com.it.operator;

import com.it.pojo.Event;
import com.it.operator.utils.SourceUtils;
import org.apache.flink.api.common.functions.Partitioner;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;

/**
 * 默认情况下调用的是rebalance方式
 */
public class Operator_Shuffle {

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment executionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment();
        executionEnvironment.setParallelism(1);
        SingleOutputStreamOperator<Event> eventSource = SourceUtils.getEventSource(executionEnvironment);
        //1.shuffle：随机分区，尽量将数据分配的均匀
        //eventSource.shuffle().print().setParallelism(3);
        //2.rebalance：轮询分区
        //eventSource.rebalance().print().setParallelism(3);
        //3.rescale：重新缩放分区，将数据分成几个小组，然后小组内进行rebalance。
        DataStreamSource<Integer> parallelSource = executionEnvironment.addSource(new RichParallelSourceFunction<Integer>() {
            @Override
            public void run(SourceContext<Integer> ctx) throws Exception {
                for (int i = 1; i <= 8; i++) {
                    //将奇/偶分别发我往不同的task子任务
                    if (i % 2 == getRuntimeContext().getIndexOfThisSubtask()) {
                        ctx.collect(i);
                    }
                }
            }

            @Override
            public void cancel() {

            }
        });
        //parallelSource.setParallelism(2).rescale().print().setParallelism(4);
        //parallelSource.setParallelism(2).global().print();
        //最后打印出的并行度也是2
        executionEnvironment.fromElements(1,2,3,4,5,6,7,8).partitionCustom(new Partitioner<Integer>() {
            @Override
            public int partition(Integer key, int numPartitions) {
                return key%2;
            }
        }, new KeySelector<Integer, Integer>() {
            @Override
            public Integer getKey(Integer value) throws Exception {
                return value;
            }
        }).print().setParallelism(4);
        executionEnvironment.execute();
    }
}
